import threading
import collections
from absl import flags

flags.DEFINE_list('windows_list',
                  '10,30,60,120,300,600,1800,3600,10800,21600',
                  'A list of windows in second')


class PriceSymbol(object):
  def __init__(self, group, exchange, currency, price_type):
    self.group = group.lower()
    self.exchange = exchange.upper()
    self.currency = currency.upper()
    self.price_type = price_type.lower()

  def __str__(self):
    return "%s.%s.%s.%s" % (self.group, self.exchange, self.currency, self.price_type)


class RingQueue(object):
  def __init__(self, size):
    self._size = size
    self._data = [None] * self._size
    self._live = [False] * self._size
    self._last_index = -1

  def _update_time(self, index):
    if index <= self._last_index:
      # Already updated. Nothing to do
      return
    prev_value = self._data[self._last_index % self._size]
    if index - self._last_index >= self._size:
      self._data = [prev_value] * self._size
      self._live = [False] * self._size
    else:
      for fill_index in range(self._last_index + 1, index + 1):
        self._data[fill_index % self._size] = prev_value
        self._live[fill_index % self._size] = False
    self._last_index = index

  def update_time(self, index):
    self._update_time(self, index)

  def update_value(self, index, value):
    if index <= self._last_index:
      # Already updated. Nothing to do
      return
    self._update_time(index - 1)
    self._data[index % self._size] = value
    self._live[index % self._size] = True
    self._last_index = index

  def _get_values(self, start_index, end_index):
    assert end_index - start_index <= self._size
    assert end_index - start_index > 0
    assert start_index >= self._last_index - self._size

    mod_start_index = start_index % self._size
    mod_end_index = end_index % self._size

    if mod_start_index < mod_end_index:
      return (self._data[mod_start_index:mod_end_index], self._live[mod_start_index:mod_end_index])
    else:
      return (self._data[mod_start_index:] + self._data[:mod_end_index],
              self._live[mod_start_index:] + self._live[:mod_end_index])

  def get_values(self, index, count):
    memory_start = self._last_index - self._size
    adj_count = min(count, self._size)
    end_index = index + count
    start_index = end_index - adj_count
    if start_index >= memory_start:
      return self._get_values(start_index, end_index)
    elif end_index > memory_start:
      data, live = self._get_values(memory_start, end_index)
      return ([None] * (memory_start - start_index) + data, [False] *
              (memory_start - start_index) + live)
    else:
      return ([None] * (end_index - start_index), [False] * (end_index - start_index))

  def get_last_index(self):
    return self._last_index

  def get_last_value(self):
    return self._data[self._last_index % self._size]


# Note that size of ring queue should be double of max MA window
class TimedRingQueue(object):
  def __init__(self, resolution_sec, window_size_sec):
    self._resolution_ns = int(resolution_sec * (10**9))
    self._window_size_ns = int(window_size_sec * (10**9))
    self._num_of_element = int(window_size_sec / resolution_sec)
    assert self._num_of_element == window_size_sec / resolution_sec
    self._queue = RingQueue(self._num_of_element)
    self._zero_idx_time_ns = 0
    self._last_fetched_time_sec = 0

  def update_time(self, fetched_time_sec):
    fetched_time_ns = int(fetched_time_sec * (10**9))
    index = int((fetched_time_ns - self._zero_idx_time_ns) / \
            self._resolution_ns)
    self._queue.update_time(index)

  def update_value(self, fetched_time_sec, value):
    fetched_time_ns = int(fetched_time_sec * (10**9))
    index = int((fetched_time_ns - self._zero_idx_time_ns) / \
            self._resolution_ns)
    self._queue.update_value(index, value)
    self._last_fetched_time_sec = fetched_time_sec

  def get_recent_sec_values(self, size_sec):
    size_ns = size_sec * 10**9
    assert size_ns <= self._windos_size_ns
    count = int(size_ns / self._resolution_ns)
    return self._queue.get_recent_values(count, count)

  def get_values(self, index, count):
    return self._queue.get_values(index, count)

  def get_last_index(self):
    return self._queue.get_last_index()

  def get_last_value(self):
    return self._queue.get_last_value()

  def get_last_fetched_time_sec(self):
    return self._last_fetched_time_sec


class MovingAverage(object):
  def __init__(self, size):
    self._size = size
    self._data = collections.deque()
    self._sum = 0
    self._cnt = 0
    self._ready = False

  @staticmethod
  def remove_none(values):
    return [value for value in values if value is not None]

  def update_values(self, values):
    if len(values) >= self._size:
      trunc_values = self.remove_none(values[-self._size:])
      self._data = collections.deque(trunc_values)
      self._sum = sum(trunc_values)
      self._cnt = len(trunc_values)
    else:
      for push_value in values:
        if push_value:
          self._data.append(push_value)
          self._sum += push_value
          self._cnt += 1
      assert self._cnt >= 0
      while len(self._data) > self._size:
        pop_value = self._data.popleft()
        if pop_value:
          self._sum -= pop_value
          self._cnt -= 1
          self._ready = True
      assert self._cnt >= 0

  def sum(self):
    return self._sum

  def average(self):
    return self._sum / (self._cnt + 1e-9)

  def ready(self):
    return self._ready


class MovingAverageCollection(object):
  def __init__(self, resolution, window_size_sec_list):
    self._mas = {}
    self._resolution = resolution
    self._window_size_sec_list = window_size_sec_list
    for second in self._window_size_sec_list:
      window_size = int(second / self._resolution)
      self._mas[window_size] = (MovingAverage(window_size), MovingAverage(window_size))
    self._last_fetched_time_sec = 0

  def update_values(self, values, live_flags, last_fetched_time_sec):
    assert isinstance(values, list)
    for ma, coverage in self._mas.values():
      ma.update_values(values)
      coverage.update_values(live_flags)
    self._last_fetched_time_sec = last_fetched_time_sec

  def get_basis(self):
    res = {}
    for window_size, ma_pair in self._mas.items():
      ma_value = ma_pair[0].average()
      coverage = ma_pair[1].average()
      res[window_size] = {
          "window_size": window_size,
          "window_size_sec": window_size * self._resolution,
          "last_fetched_time_sec": self._last_fetched_time_sec,
          "value": ma_value,
          "coverage": coverage,
          "average_interval_sec": self._resolution / (coverage + 1e-9),
          "ready": ma_pair[0].ready()
      }
    return res


class BasisStorage(object):
  def __init__(self):
    self._resolution = 0.1
    self._window_size_sec_list = \
        [int(elem) for elem in flags.FLAGS.windows_list]
    self._max_window_size_sec = max(self._window_size_sec_list) * 2
    self._products = []
    self._basis_pairs = []
    self._data = {}
    self._data_lock = {}
    self._mas = {}
    self._mas_last_index = {}
    self._tags = {}
    self._update_basis_lock = threading.Lock()
    self._add_lock = threading.Lock()

  @staticmethod
  def _calculate_basis(values_a, values_b):
    res = []
    assert len(values_a) == len(values_b)
    for value_a, value_b in zip(values_a, values_b):
      if value_a and value_b:
        res.append(value_a - value_b)
      else:
        res.append(None)
    return res

  @staticmethod
  def _logical_or(bool_a, bool_b):
    return [(elem_a or elem_b) for elem_a, elem_b in zip(bool_a, bool_b)]

  def add_product(self, price_symbol):
    assert isinstance(price_symbol, PriceSymbol)
    key = str(price_symbol)
    with self._add_lock:
      if key not in self._data:
        self._products.append(price_symbol)
        self._data[key] = TimedRingQueue(self._resolution, self._max_window_size_sec)
        self._data_lock[key] = threading.Lock()

  def add_pair(self, tag, price_symbol_a, price_symbol_b):
    assert isinstance(price_symbol_a, PriceSymbol)
    assert isinstance(price_symbol_b, PriceSymbol)
    assert price_symbol_a.group == price_symbol_b.group, \
        "Group differene between %s and %s" % (
            str(price_symbol_a), str(price_symbol_b))
    key_a = str(price_symbol_a)
    key_b = str(price_symbol_b)
    with self._add_lock:
      key_set = set(self._data.keys())
      assert key_a in key_set, (key_a, key_set)
      assert key_b in key_set, (key_b, key_set)
      if (key_a, key_b) not in self._mas:
        self._basis_pairs.append((price_symbol_a, price_symbol_b))
        self._mas[(key_a, key_b)] = MovingAverageCollection(self._resolution,
                                                            self._window_size_sec_list)
        self._mas_last_index[(key_a, key_b)] = -1
        self._tags[(key_a, key_b)] = tag

  def update_value(self, price_symbol, value, fetched_time):
    assert isinstance(price_symbol, PriceSymbol)
    key = str(price_symbol)
    with self._data_lock[key]:
      self._data[key].update_value(fetched_time, value)

  def update_basis(self, price_symbol_a, price_symbol_b):
    assert isinstance(price_symbol_a, PriceSymbol)
    assert isinstance(price_symbol_b, PriceSymbol)
    key_a = str(price_symbol_a)
    key_b = str(price_symbol_b)
    # Updating basis will be run in single thread
    with self._update_basis_lock:
      last_index_ma = self._mas_last_index[(key_a, key_b)]
      last_index_a = self._data[key_a].get_last_index()
      last_index_b = self._data[key_b].get_last_index()
      last_index = min(last_index_a, last_index_b)
      cnt = last_index - last_index_ma

      last_fetched_time_a = self._data[key_a].get_last_fetched_time_sec()
      last_fetched_time_b = self._data[key_b].get_last_fetched_time_sec()
      last_fetched_time_sec = min(last_fetched_time_a, last_fetched_time_b)

      if cnt > 0:
        with self._data_lock[key_a]:
          values_a, live_a = self._data[key_a].get_values(last_index_ma + 1, cnt)
        with self._data_lock[key_b]:
          values_b, live_b = self._data[key_b].get_values(last_index_ma + 1, cnt)
        basis_values = self._calculate_basis(values_a, values_b)
        live_flags = self._logical_or(live_a, live_b)
        self._mas[(key_a, key_b)].update_values(basis_values, live_flags, last_fetched_time_sec)
        self._mas_last_index[(key_a, key_b)] = last_index

  def get_pairs(self):
    return self._basis_pairs

  def get_basis(self):
    with self._update_basis_lock:
      basis = {}
      for price_symbol_a, price_symbol_b in self._basis_pairs:
        key_a = str(price_symbol_a)
        key_b = str(price_symbol_b)
        key = "%s,%s" % (key_a, key_b)
        basis[key] = {
            "tag": self._tags[(key_a, key_b)],
            "symbol_a": {
                "group": price_symbol_a.group,
                "exchange": price_symbol_a.exchange,
                "currency": price_symbol_a.currency,
                "price_type": price_symbol_a.price_type,
                "key": key_a
            },
            "symbol_b": {
                "group": price_symbol_b.group,
                "exchange": price_symbol_b.exchange,
                "currency": price_symbol_b.currency,
                "price_type": price_symbol_b.price_type,
                "key": key_b
            },
            "basis": self._mas[(key_a, key_b)].get_basis()
        }
      price = {}
      for price_symbol in self._products:
        key = str(price_symbol)
        price[key] = {
            "symbol": {
                "group": price_symbol.group,
                "exchange": price_symbol.exchange,
                "currency": price_symbol.currency,
                "price_type": price_symbol.price_type
            },
            "price": self._data[key].get_last_value()
        }
      return {"basis": basis, "price": price}

  def _extract_basis_group(self, value):
    try:
      if value["symbol_a"]["group"] == value["symbol_b"]["group"]:
        return value["symbol_a"]["group"]
    except Exception:
      pass
    return None

  def _extract_price_group(self, value):
    try:
      return value["symbol"]["group"]
    except Exception:
      pass
    return None

  def get_basis_per_group(self):
    whole_basis = self.get_basis()
    res = {}
    for key, value in whole_basis["basis"].items():
      group = self._extract_basis_group(value)
      if group:
        if group not in res:
          res[group] = {"basis": {}, "price": {}}
        res[group]["basis"][key] = value
    for key, value in whole_basis["price"].items():
      group = self._extract_price_group(value)
      if group:
        if group not in res:
          res[group] = {"basis": {}, "price": {}}
        res[group]["price"][key] = value
    return res
